import sys

# format:TCR\tTCRIndex\tChoBar\tchain\trawseq.    chain is V:D:J:CDR3,V contains sub class
# TCRRawSeqTableFile = sys.argv[1]
# outTable = sys.argv[2]

TCRRawSeqTableFile = "/home/hsmy/mob/p.TCRFaSeq.xls"
outTable = "/home/hsmy/mob/p.TCRExtSeq.xls"

# finish in 2020.5.30
TRABVFaFile = "/home/hsmy/code/ik/TCR/singlecell/TCRPriDesign/DB/VGeneTRAV.fa"


def getTRABVFa(inputFa, BFlag):
    # input:
    # inputFa:a file path!
    # Fafile format(>TRAV1-1\n1 AAAA)
    # output:geneFaDict

    import re
    print "from fa file get TRAV/TRBV fa..."
    geneFaDict = {}

    LineNum = -1
    GeneSeq = ""
    geneName = "NA"
    geneFaDict[geneName] = ""

    f = open(inputFa)
    line = f.readline()
    while line:
        line = line.strip("\n")
        m = re.match("(>)(\S*)", line)
        n = re.match("(\d+)(\s+)(\S+)", line)
        if m:
            geneFaDict[geneName] = GeneSeq  # save before Gene seq
            geneName = m.group(2)
            if geneName not in geneFaDict:  # a new Gene
                geneFaDict[geneName] = ""
                LineNum = 0
                GeneSeq = ""
        elif n:
            LineNum += 1
            line = line.replace(" ", "")
            line = line.replace("\t", "")
            s = re.match("(\d+)([A-Z]+)", line)
            if s:  # like:61ATCGATCG
                seq = s.group(2)
                if LineNum == 2:  # gene seq ATG start
                    if seq[0:3] == "ATG":
                        GeneSeq += seq
                    else:
                        print geneName, "warning:(this start not ATG)", line
                        GeneSeq += seq
                elif LineNum > 2:  # other line seq
                    GeneSeq += seq
            else:
                print "warning line(no num):", line
        else:
            print "this line can't see!"
        line = f.readline()
    f.close()

    if BFlag == 1:
        print "get Gene Big Class!"
        geneFaDict1 = {}
        for g in geneFaDict:
            gene = g.split("-")[0]
            if gene not in geneFaDict1:
                geneFaDict1[gene] = geneFaDict[g]

        print len(geneFaDict), len(geneFaDict1)
        print "read file end!"
        return geneFaDict1

    print "read file end!"
    return geneFaDict


def TCRSeqToAcid(tflag, TRABVFaFile, barFaDict):
    # input:
    # tflag:mouse or human
    # VGeneFaFile: a file path
    # barFaDict: bar format: bar format: TCRName\tbarcode\tV:D:J:CDR3
    # Chain type come from VGene 0:3
    # need function: getTRABVFa

    def AcidDeqExtNucSeq(AcidMatSeq, AcidTotSeq, NucTotSeq):
        Pos = AcidTotSeq.find(AcidMatSeq)
        NucMatSeq = "-"
        if Pos > -1:
            NulTotLen = len(NucTotSeq)
            AcidMatLen = len(AcidMatSeq)
            if (Pos + AcidMatLen) * 3 <= NulTotLen:
                NucMatSeq = NucTotSeq[Pos * 3:(Pos + AcidMatLen) * 3]
        return NucMatSeq

    def translateToAcid(seq):
        def StriOrderlyByNum(stri, num):
            if len(stri) % num == 0:
                return stri
            else:
                redundantNum = len(stri) % num
                totalNum = len(stri)
                return stri[0:totalNum - redundantNum]

        table = {
            'ATA': 'I', 'ATC': 'I', 'ATT': 'I', 'ATG': 'M',
            'ACA': 'T', 'ACC': 'T', 'ACG': 'T', 'ACT': 'T',
            'AAC': 'N', 'AAT': 'N', 'AAA': 'K', 'AAG': 'K',
            'AGC': 'S', 'AGT': 'S', 'AGA': 'R', 'AGG': 'R',
            'CTA': 'L', 'CTC': 'L', 'CTG': 'L', 'CTT': 'L',
            'CCA': 'P', 'CCC': 'P', 'CCG': 'P', 'CCT': 'P',
            'CAC': 'H', 'CAT': 'H', 'CAA': 'Q', 'CAG': 'Q',
            'CGA': 'R', 'CGC': 'R', 'CGG': 'R', 'CGT': 'R',
            'GTA': 'V', 'GTC': 'V', 'GTG': 'V', 'GTT': 'V',
            'GCA': 'A', 'GCC': 'A', 'GCG': 'A', 'GCT': 'A',
            'GAC': 'D', 'GAT': 'D', 'GAA': 'E', 'GAG': 'E',
            'GGA': 'G', 'GGC': 'G', 'GGG': 'G', 'GGT': 'G',
            'TCA': 'S', 'TCC': 'S', 'TCG': 'S', 'TCT': 'S',
            'TTC': 'F', 'TTT': 'F', 'TTA': 'L', 'TTG': 'L',
            'TAC': 'Y', 'TAT': 'Y', 'TAA': '_', 'TAG': '_',
            'TGC': 'C', 'TGT': 'C', 'TGA': '_', 'TGG': 'W',
        }

        seq1 = StriOrderlyByNum(seq, 3)
        protein = ""
        if len(seq1) % 3 == 0:
            for i in range(0, len(seq1), 3):
                codon = seq1[i:i + 3]
                ToAcid = table[codon]
                if ToAcid != "_":
                    protein += ToAcid
                else:
                    break
        else:
            protein = "-"
            print "please check seq num!", len(seq1) % 3

        return protein

    def ExtTwoSeqMiddle(longSeq, seq1, seq2):
        # need consider two EDL

        def getStrAllFindPos(source, key):
            PosList = []
            findPos = source.find(key)
            while findPos >= 0:
                PosList.append(findPos)
                findPos = source.find(key, findPos + 1)
            if len(PosList) > 0:
                return PosList
            else:
                return []

        PosSeq1List = getStrAllFindPos(longSeq, seq1)
        PosSeq2List = getStrAllFindPos(longSeq, seq2)

        if len(PosSeq1List) == 1 and len(PosSeq2List) == 1:
            start1 = PosSeq1List[0]
            start2 = PosSeq2List[0]
            return (longSeq[start1 + len(seq1):start2])
        elif len(PosSeq2List) > 1 and len(PosSeq1List) == 1:

            start1 = PosSeq1List[0]
            start2 = max(PosSeq2List)

            return (longSeq[start1 + len(seq1):start2])
        elif len(PosSeq1List) > 1:
            print "warning!find more than 1 CDR3"
            return ("-")
        else:
            return ("-")

    def TCRTRVseqExt(seq, kseq, CutDisMatNum):
        def getStrAllFindPos(source, key, posfrac):
            sourceLen = len(source)
            NeedBigPos = int(sourceLen * posfrac)

            PosList = []
            findPos = source.find(key)
            while findPos >= 0:
                if findPos <= NeedBigPos:
                    PosList.append(findPos)
                findPos = source.find(key, findPos + 1)
            if len(PosList) > 0:
                return PosList
            else:
                return False

        def seqmat2(seq1, seq2, CutNum):
            # input:
            # CutNum:two seq cut to the same to compare
            # output:
            # DisMatchNum(two seq dismatch num or false)
            # false(one more seq's length less than CutNum)

            L1 = len(seq1)
            L2 = len(seq2)

            if L1 < CutNum or L2 < CutNum:
                return False

            seq1 = seq1[0:CutNum]
            seq2 = seq2[0:CutNum]

            DisMatchNum = 0
            for i in range(0, CutNum):
                if seq1[i] != seq2[i]:
                    DisMatchNum += 1
            return DisMatchNum

        def getSeqMisMatchNum(seq, kseq, PosList):
            # input:
            # seq:long seq
            # kseq:known seq
            # PosList:known seq in long seq possible Pos
            # output:
            # PosMisNumDict({Pos:MisMatchNum})(0,1,2,... or "-")

            PosMisNumDict = {}
            for i in PosList:
                seqExt = seq[i:]
                MisMat = seqmat2(seqExt, kseq, len(kseq))
                if MisMat >= -1:
                    PosMisNumDict[i] = MisMat
                else:
                    PosMisNumDict[i] = "-"
            return PosMisNumDict

        def ChoMatPos(DisMatNumDict, DisMatCuf):
            # input:
            # DisMatNumDict({Pos:DisMatchNum})
            # DisMatCuf(a value of max dismatch)
            # output:
            # Pos or False
            AllMatList = []
            for i in sorted(DisMatNumDict.items(), key=lambda x: x[1], reverse=False):
                Pos = i[0]
                DisMatNum = i[1]
                if DisMatNum == 0:
                    AllMatList.append(Pos)
            # choose all match pos
            if len(AllMatList) > 0:
                AllMatList.sort()
                return AllMatList[0]
            # dismatch num is more low and pos is Front
            else:
                for i in sorted(DisMatNumDict.items(), key=lambda x: x[1], reverse=False):
                    Pos = i[0]
                    DisMatNum = i[1]

                    if DisMatNum <= DisMatCuf:
                        return Pos
                return False

        # TCRTRVseqExt main
        # ATG
        HeadKSeq = kseq[0:3]
        PosList = getStrAllFindPos(seq, HeadKSeq, 0.75)

        if PosList != False:
            DisMatNumDict = getSeqMisMatchNum(seq, kseq, PosList)
            FinalPos = ChoMatPos(DisMatNumDict, CutDisMatNum)
            return seq[FinalPos:]
        else:
            return "-"

    # base function end!

    # TCRSeqToAcid main
    # read gene - Fa Dict
    VGeneFaDict = getTRABVFa(TRABVFaFile, 1)
    VGeneFaDict2 = getTRABVFa(TRABVFaFile, 0)

    # extract TCR seq and translate to Acid seq
    barFaExtDict = {}
    barAcidDict = {}
    barLinkAcidDict = {}
    barCDR3SeqDict = {}
    VGeneBarListDict = {}

    # PseudoGeneList = ["TRAV11", "TRAV15", "TRAV28", "TRAV31", "TRAV32", "TRAV33", "TRAV37", "TRBV1", "TRBV5-2",
    #                   "TRBV5-3", "TRBV6-2", "TRBV7-1", "TRBV7-4", "TRBV7-5", "TRBV8-1", "TRBV8-2", "TRBV12-1",
    #                   "TRBV12-2", "TRBV17", "TRBV21", "TRBV22", "TRBV23", "TRBV26"]
    PseudoGeneList = []
    for i in barFaDict:
        TempBarName = i.split("\t")
        Barcode = TempBarName[1]
        # TRAV1-1
        VGene = TempBarName[2].split(":")[0].replace("/", "")
        # TRAV1
        VGene2 = VGene.split("-")[0]
        CDR3 = TempBarName[2].split(":")[3]
        chain = VGene[0:3]

        if VGene in VGeneBarListDict:
            if Barcode not in VGeneBarListDict[VGene]:
                VGeneBarListDict[VGene].append(Barcode)
        else:
            VGeneBarListDict[VGene] = []
            VGeneBarListDict[VGene].append(Barcode)

        SeqFaExt = "-"
        SeqFaAcid = "-"
        JCLink = "-"
        CDR3Seq = "-"

        if tflag != "Mouse":
            flag = 1
        else:
            flag = 0

        if VGene in PseudoGeneList:
            print("Pseudo Gene:", i)
            flag = 0
        elif VGene2 not in VGeneFaDict:
            print("not in VGeneFaDict:", i)
            flag = 0

        # get acid seq
        if flag == 1:
            # choose Gene Seq
            if VGene in VGeneFaDict2:
                StandSeq = VGeneFaDict2[VGene]
            else:
                StandSeq = VGeneFaDict[VGene2]
            SeqFaExt = TCRTRVseqExt(barFaDict[i], StandSeq[0:20], 5)

            if SeqFaExt == "-":
                print(i, VGene)
            else:
                # can extact fa seq and get acid seq and get link acid seq
                endStri = ""
                if chain == "TRA":
                    endStri = "IQN"
                elif chain == "TRB":
                    endStri = "EDL"
                else:
                    print("chain is error!", chain)

                SeqFaAcid = translateToAcid(SeqFaExt)
                if endStri != "":
                    JCLink = ExtTwoSeqMiddle(SeqFaAcid, CDR3, endStri)
                    JCLink = "%s%s" % (JCLink, endStri)

        # get CDR3 Nucleic seq
        if SeqFaExt != "-" and SeqFaAcid != "-":
            CDR3Seq = AcidDeqExtNucSeq(CDR3, SeqFaAcid, SeqFaExt)

        if SeqFaExt[0:3] != "ATG":
            # print "not from ATG",SeqFaExt
            SeqFaExt = "-"
            SeqFaAcid = "-"
            JCLink = "-"
            CDR3Seq = "-"

        if JCLink == "-EDL" or JCLink == "-IQN":
            JCLink = "-"

        # key : barcode\tTRX\tTRXV:CDR3
        barAcidDict[i] = SeqFaAcid
        barLinkAcidDict[i] = JCLink
        barFaExtDict[i] = SeqFaExt
        barCDR3SeqDict[i] = CDR3Seq

    barSeqResult = [barFaExtDict, barAcidDict, barLinkAcidDict, barCDR3SeqDict, VGeneBarListDict]
    return barSeqResult
    # TCRSeqToAcid main end.


# main
BarFaKeyList = []
BarKeyLineDict = {}
BarFaDict = {}

f = open(TCRRawSeqTableFile)
f.readline()
line = f.readline()
while line:
    temp = line.strip("\n").split("\t")

    key = "%s\t%s\t%s" % (temp[0], temp[2], temp[3])
    TCRSeq = temp[4]

    if key not in BarFaDict:
        BarFaDict[key] = TCRSeq
        BarFaKeyList.append(key)
        BarKeyLineDict[key] = line.strip("\n")

    line = f.readline()
f.close()

TCRSeqResult = TCRSeqToAcid("human", TRABVFaFile, BarFaDict)
barFaExtDict = TCRSeqResult[0]
barAcidDict = TCRSeqResult[1]
barLinkAcidDict = TCRSeqResult[2]
barCDR3SeqDict = TCRSeqResult[3]
VGeneBarListDict = TCRSeqResult[4]

NewLineList = []
Head = "TCR\tTCRIndex\tChoBar\tchain\trawseq\tVGene\tDGene\tJGene\tCDR3Acid1\ttype\tExtSeq\tAcid\tLinkAcid\tCDR3Acid2\tCDR3Dna1\n"
NewLineList.append(Head)
for i in BarFaKeyList:
    temp = i.split("\t")
    TCRChainTemp = temp[2].split(":")
    VGene = TCRChainTemp[0]
    DGene = TCRChainTemp[1]
    JGene = TCRChainTemp[2]
    CDR3Acid1 = TCRChainTemp[3]

    if barLinkAcidDict[i] != "-":
        CDR3Acid2 = CDR3Acid1 + barLinkAcidDict[i]
    else:
        CDR3Acid2 = "-"

    ctype = "-"
    if VGene[0:3] == "TRA":
        ctype = "A"
    elif VGene[0:3] == "TRB":
        ctype = "B"

    newline = "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n" % (
        BarKeyLineDict[i],
        VGene, DGene, JGene, CDR3Acid1, ctype,
        barFaExtDict[i], barAcidDict[i], barLinkAcidDict[i],
        CDR3Acid2, barCDR3SeqDict[i]
    )

    NewLineList.append(newline)

w = file(outTable, "w+")
for i in NewLineList:
    w.write(i)
w.close()
